import logging
from psycopg2.extensions import connection as _connection
from psycopg2.extras import DictConnection, DictCursor
from twisted.internet.defer import inlineCallbacks, returnValue
from twisted.python.util import mergeFunctionMetadata
from txpostgres.txpostgres import ConnectionPool

from littlesite import config


pool = ConnectionPool(None)

class CustomCursor(DictCursor):
    def execute(self, query, vars=None):
        try:
            return DictCursor.execute(self, query, vars)
        finally:
            if vars:
                logging.debug('DB executing %r with %r' % (query, vars))
            else:
                logging.debug('DB executing %r' % query)


class CustomConnection(DictConnection):
    def cursor(self, name=None):
        if name is None:
            return _connection.cursor(self, cursor_factory=CustomCursor)
        else:
            return _connection.cursor(self, name, cursor_factory=CustomCursor)


def initConnectionPool():
    global pool
    pool = ConnectionPool(None, database=config.db_name,
                          connection_factory=CustomConnection)
    d = pool.start()
    d.consumeErrors = True
    return d


currentTransaction = None

def getCurrentTransaction():
    assert currentTransaction, 'No transaction available!'
    return currentTransaction


class Transaction:
    def __init__(self, cursor):
        self._cursor = cursor

    @inlineCallbacks
    def update(self, sql, params=None):
        yield self._cursor.execute(sql, params)

    @inlineCallbacks
    def select(self, sql, params=None):
        yield self._cursor.execute(sql, params)
        returnValue([dict(row) for row in self._cursor.fetchall()])


def _doInTransaction(f, args, kwargs):
    def unwindCursor(cursor):
        global currentTransaction
        try:
            currentTransaction = Transaction(cursor)
            return f(*args, **kwargs)
        finally:
            currentTransaction = None
    return pool.runInteraction(unwindCursor)

def doInTransaction(f):
    logging.debug('Wrapping %s with doInTransaction' % f.func_name)
    f_ = inlineCallbacks(f)
    def unwindGenerator(*args, **kwargs):
        return _doInTransaction(f_, args, kwargs)
    return mergeFunctionMetadata(f, unwindGenerator)
